# data_storage/database.py - Database for storing event information

import os
import sqlite3
import json
import logging
from datetime import datetime
import threading

logger = logging.getLogger(__name__)


class Database:
    """
    SQLite database for storing video event information
    """

    def __init__(self, db_path="events.db"):
        self.db_path = db_path
        self.conn = None
        self.lock = threading.RLock()  # 添加可重入锁

        # Initialize the database
        self.initialize()

    def initialize(self):
        """Initialize the database and create tables if they don't exist"""
        try:
            # Create directory if it doesn't exist
            db_dir = os.path.dirname(self.db_path)
            if db_dir and not os.path.exists(db_dir):
                os.makedirs(db_dir)

            with self.lock:  # 使用锁确保线程安全
                # 关闭可能存在的旧连接
                if self.conn:
                    try:
                        self.conn.close()
                    except:
                        pass

                # 使用 check_same_thread=False 允许在不同线程中使用
                self.conn = sqlite3.connect(self.db_path, check_same_thread=False)

                # Create tables
                self._create_tables()

            logger.info(f"Database initialized: {self.db_path}")

        except Exception as e:
            logger.error(f"Error initializing database: {str(e)}")

    def _create_tables(self):
        """Create the necessary tables in the database"""
        cursor = self.conn.cursor()

        # Videos table
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS videos (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            path TEXT UNIQUE,
            duration REAL,
            frame_count INTEGER,
            width INTEGER,
            height INTEGER,
            fps REAL,
            import_date TEXT
        )
        ''')

        # Events table
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS events (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            video_id INTEGER,
            type TEXT,
            start_time REAL,
            end_time REAL,
            duration REAL,
            confidence REAL,
            frame_start INTEGER,
            frame_end INTEGER,
            output_file TEXT,
            metadata TEXT,
            created_at TEXT,
            modified_at TEXT,
            FOREIGN KEY (video_id) REFERENCES videos(id)
        )
        ''')

        self.conn.commit()

    def insert_video(self, video_info):
        """
        Insert a video into the database

        Args:
            video_info: Dictionary with video information

        Returns:
            ID of the inserted video
        """
        if not self.conn:
            self.initialize()

        try:
            cursor = self.conn.cursor()

            # Check if video already exists
            cursor.execute(
                "SELECT id FROM videos WHERE path = ?",
                (video_info.get('path'),)
            )
            existing = cursor.fetchone()

            if existing:
                return existing[0]

            # Insert the video
            cursor.execute('''
            INSERT INTO videos (
                path, duration, frame_count, width, height, fps, import_date
            ) VALUES (?, ?, ?, ?, ?, ?, ?)
            ''', (
                video_info.get('path', ''),
                video_info.get('duration', 0),
                video_info.get('frame_count', 0),
                video_info.get('width', 0),
                video_info.get('height', 0),
                video_info.get('fps', 0),
                datetime.now().isoformat()
            ))

            self.conn.commit()
            return cursor.lastrowid

        except Exception as e:
            logger.error(f"Error inserting video: {str(e)}")
            return None

    def insert_event(self, event):
        """
        Insert an event into the database

        Args:
            event: Dictionary with event information

        Returns:
            ID of the inserted event
        """
        if not self.conn:
            self.initialize()

        try:
            cursor = self.conn.cursor()

            # Get video ID
            video_path = event.get('video_path', '')
            video_id = None

            cursor.execute("SELECT id FROM videos WHERE path = ?", (video_path,))
            video_result = cursor.fetchone()

            if video_result:
                video_id = video_result[0]
            else:
                # Insert the video if it doesn't exist
                cursor.execute('''
                INSERT INTO videos (path, import_date) 
                VALUES (?, ?)
                ''', (video_path, datetime.now().isoformat()))
                video_id = cursor.lastrowid

            # Extract metadata
            metadata = {k: v for k, v in event.items()
                        if k not in ['id', 'video_path', 'type', 'start_time',
                                     'end_time', 'duration', 'confidence',
                                     'frames', 'output_file']}

            # Insert the event
            cursor.execute('''
            INSERT INTO events (
                video_id, type, start_time, end_time, duration, confidence,
                frame_start, frame_end, output_file, metadata, created_at, modified_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ''', (
                video_id,
                event.get('type', ''),
                event.get('start_time', 0),
                event.get('end_time', 0),
                event.get('duration', 0),
                event.get('confidence', 0),
                event.get('frames', {}).get('start', 0),
                event.get('frames', {}).get('end', 0),
                event.get('output_file', ''),
                json.dumps(metadata),
                datetime.now().isoformat(),
                datetime.now().isoformat()
            ))

            self.conn.commit()
            event_id = cursor.lastrowid

            return event_id

        except Exception as e:
            logger.error(f"Error inserting event: {str(e)}")
            return None

    def update_event(self, event):
        """
        Update an event in the database

        Args:
            event: Dictionary with event information

        Returns:
            True if successful, False otherwise
        """
        if not self.conn:
            self.initialize()

        try:
            with self.lock:  # 使用锁确保线程安全
                cursor = self.conn.cursor()

                # Extract metadata
                metadata = {k: v for k, v in event.items()
                            if k not in ['id', 'video_path', 'type', 'start_time',
                                         'end_time', 'duration', 'confidence',
                                         'frames', 'output_file']}

                # Update the event
                cursor.execute('''
                UPDATE events
                SET type = ?, start_time = ?, end_time = ?, duration = ?, confidence = ?,
                    frame_start = ?, frame_end = ?, output_file = ?, metadata = ?, modified_at = ?
                WHERE id = ?
                ''', (
                    event.get('type', ''),
                    event.get('start_time', 0),
                    event.get('end_time', 0),
                    event.get('duration', 0),
                    event.get('confidence', 0),
                    event.get('frames', {}).get('start', 0),
                    event.get('frames', {}).get('end', 0),
                    event.get('output_file', ''),
                    json.dumps(metadata),
                    datetime.now().isoformat(),
                    event.get('id')
                ))

                self.conn.commit()
                return cursor.rowcount > 0

        except Exception as e:
            logger.error(f"Error updating event: {str(e)}")
            return False

    def delete_event(self, event_id):
        """
        Delete an event from the database

        Args:
            event_id: ID of the event to delete

        Returns:
            True if successful, False otherwise
        """
        if not self.conn:
            self.initialize()

        try:
            cursor = self.conn.cursor()

            # Delete the event
            cursor.execute("DELETE FROM events WHERE id = ?", (event_id,))

            self.conn.commit()
            return cursor.rowcount > 0

        except Exception as e:
            logger.error(f"Error deleting event: {str(e)}")
            return False

    def get_event(self, event_id):
        """
        Get an event from the database

        Args:
            event_id: ID of the event to retrieve

        Returns:
            Event dictionary
        """
        if not self.conn:
            self.initialize()

        try:
            cursor = self.conn.cursor()

            # Get the event
            cursor.execute('''
            SELECT e.id, e.type, e.start_time, e.end_time, e.duration, e.confidence,
                   e.frame_start, e.frame_end, e.output_file, e.metadata, v.path
            FROM events e
            JOIN videos v ON e.video_id = v.id
            WHERE e.id = ?
            ''', (event_id,))

            row = cursor.fetchone()

            if row:
                event = {
                    'id': row[0],
                    'type': row[1],
                    'start_time': row[2],
                    'end_time': row[3],
                    'duration': row[4],
                    'confidence': row[5],
                    'frames': {
                        'start': row[6],
                        'end': row[7]
                    },
                    'output_file': row[8],
                    'video_path': row[10]
                }

                # Parse metadata
                try:
                    metadata = json.loads(row[9])
                    event.update(metadata)
                except:
                    pass

                return event

            return None

        except Exception as e:
            logger.error(f"Error getting event: {str(e)}")
            return None

    def get_all_events(self, event_type=None, video_id=None):
        """
        Get all events from the database

        Args:
            event_type: Filter by event type
            video_id: Filter by video ID

        Returns:
            List of event dictionaries
        """
        if not self.conn:
            self.initialize()

        try:
            cursor = self.conn.cursor()

            # Build the query
            query = '''
            SELECT e.id, e.type, e.start_time, e.end_time, e.duration, e.confidence,
                   e.frame_start, e.frame_end, e.output_file, e.metadata, v.path
            FROM events e
            JOIN videos v ON e.video_id = v.id
            '''

            params = []

            if event_type or video_id:
                query += " WHERE"

                if event_type:
                    query += " e.type = ?"
                    params.append(event_type)

                    if video_id:
                        query += " AND e.video_id = ?"
                        params.append(video_id)
                elif video_id:
                    query += " e.video_id = ?"
                    params.append(video_id)

            query += " ORDER BY e.start_time"

            # Execute the query
            cursor.execute(query, params)

            rows = cursor.fetchall()
            events = []

            for row in rows:
                event = {
                    'id': row[0],
                    'type': row[1],
                    'start_time': row[2],
                    'end_time': row[3],
                    'duration': row[4],
                    'confidence': row[5],
                    'frames': {
                        'start': row[6],
                        'end': row[7]
                    },
                    'output_file': row[8],
                    'video_path': row[10]
                }

                # Parse metadata
                try:
                    metadata = json.loads(row[9])
                    event.update(metadata)
                except:
                    pass

                events.append(event)

            return events

        except Exception as e:
            logger.error(f"Error getting events: {str(e)}")
            return []

    def get_video_info(self, video_id):
        """
        Get video information from the database

        Args:
            video_id: ID of the video

        Returns:
            Video information dictionary
        """
        if not self.conn:
            self.initialize()

        try:
            cursor = self.conn.cursor()

            # Get the video
            cursor.execute('''
            SELECT id, path, duration, frame_count, width, height, fps, import_date
            FROM videos
            WHERE id = ?
            ''', (video_id,))

            row = cursor.fetchone()

            if row:
                return {
                    'id': row[0],
                    'path': row[1],
                    'duration': row[2],
                    'frame_count': row[3],
                    'width': row[4],
                    'height': row[5],
                    'fps': row[6],
                    'import_date': row[7]
                }

            return None

        except Exception as e:
            logger.error(f"Error getting video info: {str(e)}")
            return None

    def get_video_by_path(self, video_path):
        """
        Get video information by path

        Args:
            video_path: Path to the video file

        Returns:
            Video information dictionary
        """
        if not self.conn:
            self.initialize()

        try:
            cursor = self.conn.cursor()

            # Get the video
            cursor.execute('''
            SELECT id, path, duration, frame_count, width, height, fps, import_date
            FROM videos
            WHERE path = ?
            ''', (video_path,))

            row = cursor.fetchone()

            if row:
                return {
                    'id': row[0],
                    'path': row[1],
                    'duration': row[2],
                    'frame_count': row[3],
                    'width': row[4],
                    'height': row[5],
                    'fps': row[6],
                    'import_date': row[7]
                }

            return None

        except Exception as e:
            logger.error(f"Error getting video by path: {str(e)}")
            return None

    def get_all_videos(self):
        """
        Get all videos from the database

        Returns:
            List of video information dictionaries
        """
        if not self.conn:
            self.initialize()

        try:
            cursor = self.conn.cursor()

            # Get all videos
            cursor.execute('''
            SELECT id, path, duration, frame_count, width, height, fps, import_date
            FROM videos
            ORDER BY import_date DESC
            ''')

            rows = cursor.fetchall()
            videos = []

            for row in rows:
                videos.append({
                    'id': row[0],
                    'path': row[1],
                    'duration': row[2],
                    'frame_count': row[3],
                    'width': row[4],
                    'height': row[5],
                    'fps': row[6],
                    'import_date': row[7]
                })

            return videos

        except Exception as e:
            logger.error(f"Error getting all videos: {str(e)}")
            return []

    def get_event_count_by_type(self):
        """
        Get count of events by type

        Returns:
            Dictionary of event types and counts
        """
        if not self.conn:
            self.initialize()

        try:
            cursor = self.conn.cursor()

            # Get event counts
            cursor.execute('''
            SELECT type, COUNT(*)
            FROM events
            GROUP BY type
            ''')

            rows = cursor.fetchall()
            counts = {}

            for row in rows:
                counts[row[0]] = row[1]

            return counts

        except Exception as e:
            logger.error(f"Error getting event counts: {str(e)}")
            return {}

    def clear_events(self, video_id=None):
        """
        Clear events from the database

        Args:
            video_id: ID of the video (if None, clears all events)

        Returns:
            Number of events deleted
        """
        if not self.conn:
            self.initialize()

        try:
            cursor = self.conn.cursor()

            if video_id:
                # Delete events for the specified video
                cursor.execute("DELETE FROM events WHERE video_id = ?", (video_id,))
            else:
                # Delete all events
                cursor.execute("DELETE FROM events")

            self.conn.commit()
            return cursor.rowcount

        except Exception as e:
            logger.error(f"Error clearing events: {str(e)}")
            return 0

    def get_event_statistics(self):
        """
        Get statistics about events

        Returns:
            Statistics dictionary
        """
        if not self.conn:
            self.initialize()

        try:
            cursor = self.conn.cursor()
            stats = {}

            # Get total count
            cursor.execute("SELECT COUNT(*) FROM events")
            stats['total_events'] = cursor.fetchone()[0]

            # Get count by type
            stats['events_by_type'] = self.get_event_count_by_type()

            # Get average duration
            cursor.execute("SELECT AVG(duration) FROM events")
            stats['avg_duration'] = cursor.fetchone()[0] or 0

            # Get average confidence
            cursor.execute("SELECT AVG(confidence) FROM events")
            stats['avg_confidence'] = cursor.fetchone()[0] or 0

            # Get video count
            cursor.execute("SELECT COUNT(*) FROM videos")
            stats['total_videos'] = cursor.fetchone()[0]

            return stats

        except Exception as e:
            logger.error(f"Error getting event statistics: {str(e)}")
            return {}

    def close(self):
        """Close the database connection"""
        if self.conn:
            self.conn.close()
            self.conn = None